import os
import whisper

class WhisperBatchTranscriber:
    """
    批量音频转写为字幕（SRT）文件的工具类。
    支持设置模型、语言、输入输出路径等。
    """

    def __init__(self, input_folder=None, output_folder=None,
                 model_name="base", language="ja", device="cpu"):
        """
        初始化转录器。

        Args:
            input_folder (str): 输入音频文件夹路径
            output_folder (str): 输出字幕文件夹路径
            model_name (str): Whisper 模型名称（tiny, base, small, medium, large）
            language (str): 识别语言代码，如 'ja'（日语）、'zh'（中文）、'en'（英语）
            device (str): 运行设备，'cpu' 或 'cuda'
        """
        self.input_folder = input_folder
        self.output_folder = output_folder
        self.model_name = model_name
        self.language = language
        self.device = device
        self.model = None  # 延迟加载模型

    # ===== Setter 方法 =====
    def set_input_folder(self, path: str):
        """设置输入音频文件夹"""
        self.input_folder = path

    def set_output_folder(self, path: str):
        """设置输出字幕文件夹"""
        self.output_folder = path

    def set_model_name(self, name: str):
        """设置 Whisper 模型名称"""
        self.model_name = name

    def set_language(self, lang: str):
        """设置识别语言（ISO 639-1 代码，如 'ja', 'zh', 'en'）"""
        self.language = lang

    def set_device(self, device: str):
        """设置运行设备（'cpu' 或 'cuda'）"""
        self.device = device

    # ===== 内部辅助函数 =====
    def _format_time(self, seconds: float) -> str:
        """将秒数转换为 SRT 时间格式 00:00:00,000"""
        millisec = int((seconds - int(seconds)) * 1000)
        h = int(seconds // 3600)
        m = int((seconds % 3600) // 60)
        s = int(seconds % 60)
        return f"{h:02d}:{m:02d}:{s:02d},{millisec:03d}"

    # ===== 主功能函数 =====
    def transcribe_all(self):
        """批量识别输入文件夹内的音频文件并输出 SRT 文件"""
        if not self.input_folder or not os.path.isdir(self.input_folder):
            raise ValueError("❌ 无效的输入文件夹，请先设置 input_folder。")

        os.makedirs(self.output_folder, exist_ok=True)

        # 加载模型（延迟加载）
        if self.model is None:
            print(f"🔍 正在加载 Whisper 模型：{self.model_name} ({self.device}) ...")
            self.model = whisper.load_model(self.model_name, device=self.device)
            print("✅ 模型加载完成。")

        # 遍历音频文件
        for filename in os.listdir(self.input_folder):
            if not filename.lower().endswith((".mp3", ".wav", ".m4a", ".flac")):
                continue

            input_path = os.path.join(self.input_folder, filename)
            output_path = os.path.join(self.output_folder, os.path.splitext(filename)[0] + ".srt")

            # 跳过已存在的文件
            if os.path.exists(output_path):
                print(f"⏩ 跳过（已存在）：{output_path}")
                continue

            print(f"🎧 正在识别：{filename} ...")

            # 执行语音识别
            result = self.model.transcribe(input_path, language=self.language)

            # 生成 SRT 内容
            srt_lines = []
            for i, seg in enumerate(result["segments"], start=1):
                start_time = self._format_time(seg["start"])
                end_time = self._format_time(seg["end"])
                text = seg["text"].strip()

                srt_lines.append(f"{i}")
                srt_lines.append(f"{start_time} --> {end_time}")
                srt_lines.append(text)
                srt_lines.append("")

            # 写出 SRT 文件
            with open(output_path, "w", encoding="utf-8") as f:
                f.write("\n".join(srt_lines))

            print(f"✅ 已生成字幕：{output_path}")

        print("🎉 全部音频处理完成！")

# ================= 使用示例 =================
if __name__ == "__main__":
    transcriber = WhisperBatchTranscriber(
        input_folder=r"C:\Users\33720\Desktop\xmly27\Funingna",
        output_folder=r"C:\Users\33720\Desktop\xmly27\srt",
        model_name="medium",
        language="ja",
        device="cpu"  # 如果有显卡，可改为 "cuda"
    )

    # 你也可以用 setter 修改参数
    # transcriber.set_model_name("small")
    # transcriber.set_language("zh")

    transcriber.transcribe_all()
